import os
import matplotlib.pyplot as plt
import torch

def visualize_results(masked_image, mask, gt_image, output, step, output_dir, prefix="train"):
    """可视化修复结果并保存"""
    plt.figure(figsize=(15, 5))
    
    # 将张量转换为numpy数组并调整范围
    def to_display(img):
        return (img.detach().cpu().permute(1, 2, 0).numpy() + 1) / 2
    
    # 显示带掩码图像
    plt.subplot(1, 4, 1)
    plt.imshow(to_display(masked_image))
    plt.title("Masked Image")
    plt.axis('off')
    
    # 显示掩码
    plt.subplot(1, 4, 2)
    plt.imshow(mask.detach().cpu().squeeze().numpy(), cmap='gray')
    plt.title("Mask")
    plt.axis('off')
    
    # 显示修复结果
    plt.subplot(1, 4, 3)
    plt.imshow(to_display(output))
    plt.title("Inpainting Result")
    plt.axis('off')
    
    # 显示原始图像
    plt.subplot(1, 4, 4)
    plt.imshow(to_display(gt_image))
    plt.title("Ground Truth")
    plt.axis('off')
    
    # 保存图像
    vis_path = os.path.join(output_dir, f"visualizations/{prefix}")
    os.makedirs(vis_path, exist_ok=True)
    plt.savefig(os.path.join(vis_path, f"step_{step}.png"))
    plt.close()